from flask import Flask, request, jsonify
from flask_cors import CORS
import json
import os
import base64
from PIL import Image
from io import BytesIO
import shutil

app = Flask(__name__)
app.secret_key = b'_5#y3332323xc'
CORS(app, origins='*')

log_save_path = 'runs/detect_temp_logs'
detect_file_path = 'runs/detect/exp-video'
detect_txt_path = detect_file_path + '/labels'
img_save_path = './runs/detect_temp_images'
img_size = (640, 640)
detect_cmd_prefix = 'python ./detect.py'
detect_cmd_source = '--source ' + img_save_path
detect_cmd_weights = '--weights ./runs/train/bike-0804/weights/best.pt'
detect_cmd_end = '--save-txt'


def save_base64_image(base64_str, file_path):
    img_data = base64.b64decode(base64_str)
    img = Image.open(BytesIO(img_data)).resize(img_size, Image.Resampling.LANCZOS)
    img.save(file_path)


def get_detections():
    all_detections = []
    for filename in os.listdir(detect_file_path):
        # 读取 Yolo 文件
        detections = []
        if filename == 'labels':
            continue
        if os.path.exists(os.path.join(detect_txt_path, filename.split('.')[0] + '.txt')):
            with open(os.path.join(detect_txt_path, filename.split('.')[0] + '.txt'), 'r') as f:
                for line in f:
                    class_id, x, y, width, height = map(float, line.strip().split())
                    detections.append({
                        'class_id': int(class_id),
                        'x': x,
                        'y': y,
                        'width': width,
                        'height': height
                    })
        # 读取图像文件并转换为 base64 编码
        with open(os.path.join(detect_file_path, filename), 'rb') as f:
            image_data = f.read()
            image_base64 = base64.b64encode(image_data).decode('utf-8')
            # 获取图片后缀名，添加类型（先输出看下）
            image_base64_prefix = "data:image/" + filename.split('.')[-1] + ";base64,"
            # 组装结果
            result = {
                'detections': detections,
                'image_name': filename,
                'image_base64': image_base64_prefix + image_base64
            }
            all_detections.append(result)

    # 读取log日志文件
    logs = []
    with open(log_save_path + '/log.txt', 'r') as f:
        # 读取所有行
        lines = f.readlines()
        # 打印所有行
        for line in lines:
            logs.append(line)

    # 组装最终json结果
    final_res = {
        'logs': logs,
        'details': all_detections
    }

    # 清除目录
    if os.path.exists(img_save_path):
        shutil.rmtree(img_save_path)
    if os.path.exists(detect_file_path):
        shutil.rmtree(detect_file_path)
    if os.path.exists(log_save_path):
        shutil.rmtree(log_save_path)

    return final_res


@app.route('/image-yolov5-detect', methods=['POST'])
def detect_image_obj():
    try:
        # 获取传入的参数
        data = json.loads(request.get_data())
        image_codes = data.get('imageCodes')
        # 处理图片
        if not os.path.exists(os.path.join(img_save_path)):
            os.makedirs(os.path.join(img_save_path))
        for index, image_code in enumerate(image_codes):
            # 判断图片类型
            img_type, img_code = image_code.split(',')[0], image_code.split(',')[1]
            start = img_type.find('/') + 1
            end = img_type.find(';')
            mime_type = img_type[start:end]
            if mime_type not in ['png', 'jpg', 'jpeg']:
                return jsonify({"code": 500, "msg": "图片类型错误，必须为png或jpg!"})
            save_base64_image(img_code, os.path.join(img_save_path, str(index) + '.' + mime_type))
        # 执行预测命令
        cmd = detect_cmd_prefix + ' ' + detect_cmd_source + ' ' + detect_cmd_weights + ' ' + detect_cmd_end
        print('正在执行预测脚本：', cmd)
        os.system(cmd)
        # 整理输出结果并删除目录
        return jsonify(get_detections())


    # 异常参数处理
    except Exception:
        print('输入参数异常:', Exception)
        return jsonify({'code': '500', 'msg': '输入参数异常'}, ensure_ascii=False)


if __name__ == '__main__':
    os.chdir(os.path.abspath(os.path.dirname(__file__)))
    app.run(debug=True, host='0.0.0.0', port='9528')
